import numpy as np
import matplotlib.pyplot as plt
import random
from PIL import Image

# 初始化参数
N = 400  # 系统大小
n_iter = 1000  # 迭代次数
P_num = 1  # 规则1中的随机概率
numbers = [0, 1, -1]#偏移量
state = 180 #晶粒取向角度数量

# 定义元胞自动机模型
def CA_model(N, n_iter):
    # 创建空的晶粒图像
    grains = np.zeros((N, N), dtype=int)
    # 生成初始晶粒
    for i in range(N):
        for j in range(N):
            # 晶粒有不同取向
            grains[i, j] = random.randint(0, state)
    # 迭代模拟晶粒生长和退化过程
#    grainstest = [[1,1,0],[1,0,0],[0,0,0]]
#   res = find_neighbors2(grainstest,1,1)
#   print(res)

    for t in range(n_iter):  # t = 0 ~ 999
        print("更新ing" + str(t))

        # 晶界迁移规则
        for i in range(N):
            for j in range(N):
                if (random.random() < P_num):  # 规则1 i--> (i+N)%N     j---> (j+N)%N
                    num = find_neighbors2(grains, i, j)  # 规则2
                    if (num != -1 and grains[i,j] != num):
                        grains[i, j] = num
                    elif (find_neighbors3(grains, i, j) != -1 and grains[i,j] != find_neighbors3(grains, i, j)):  # 规则3
                        grains[i, j] = find_neighbors3(grains, i, j)
                    elif (find_neighbors4(grains, i, j) != -1 and grains[i,j] != find_neighbors4(grains, i, j)):  # 规则4
                        grains[i, j] = find_neighbors4(grains, i, j)
                    else:
                        #print("success-----------------------------------")
                        irandom = random.choice(numbers)
                        jrandom = random.choice(numbers)
                        if ((gbe(grains, i, j) - gbe(grains, (i + irandom + N) % N, (j + jrandom+ N) % N)) > 0):  # 规则5
                            grains[i, j] = grains[(i + irandom + N) % N, (j + jrandom+ N) % N]
        if ( t % 50 == 0):
            plot_array(grains,t)
            grain_size = count_grain_size(grains)
            plot_grain_size_distribution(grain_size)


# 二维矩阵显示为图片
def plot_array(arr,t):
    cmap = plt.get_cmap('viridis', np.max(arr) - np.min(arr) + 1)
    plt.imshow(arr, cmap=cmap, interpolation='nearest')
    plt.savefig('./img/grains' + str(t) + '.png')
    #plt.show()



def visualize_array(arr):
    # 定义颜色对应的RGB值
    colors = [(255, 255, 255), (255, 0, 0), (0, 255, 0), (0, 0, 255)]
    # 创建一个空白的图片对象
    img = Image.new('RGB', (len(arr[0]), len(arr)))
    # 遍历二维数组，将每个数值映射到对应的颜色并填充到图片中
    for i in range(len(arr)):
        for j in range(len(arr[0])):
            color_index = min(arr[i][j] // 60, 3)
            img.putpixel((j, i), colors[color_index])
    # 显示图片
    img.show()

# 规则2
def find_neighbors2(matrix, row, col):
    # 获取二维数组的行数和列数
    num_rows = len(matrix)
    num_cols = len(matrix[0])

    # 定义邻近位置的偏移量
    offsets = [(-1, -1), (-1, 0), (-1, 1),
               (0, -1), (0, 1),
               (1, -1), (1, 0), (1, 1)]

    # 统计邻近位置中每个数字的出现次数
    count = {}
    for offset in offsets:
        # 计算邻近位置的行列坐标
        i = (row + offset[0]) % num_rows
        j = (col + offset[1]) % num_cols
        # 获取邻近位置的值
        value = matrix[i][j]
        # 统计出现次数
        count[value] = count.get(value, 0) + 1

    # 判断是否存在5个及以上的相同值
    for key, value in count.items():
        if key != -1 and value >= 5:
            return key

    # 如果不存在符合条件的数值，则返回-1
    return -1



# 规则3
def find_neighbors3(matrix, row, col):
    m, n = len(matrix), len(matrix[0])
    count = 0
    for i in range(row - 1, row + 2):
        for j in range(col - 1, col + 2):
            if i == row and j == col:
                continue
            if matrix[i % m][j % n] == matrix[row][col]:
                count += 1
    if count >= 3:
        return matrix[row][col]
    else:
        return -1


# 规则4
def find_neighbors4(matrix, row, col):
    m, n = len(matrix), len(matrix[0])
    count = 0
    value = matrix[row][col]
    for i, j in [(row - 1, col - 1), (row - 1, col + 1), (row + 1, col - 1), (row + 1, col + 1)]:
        if matrix[i % m][j % n] == value:
            count += 1
    if count >= 3:
        return value
    else:
        return -1


# 界面能计算
def gbe(arr, i, j):
    n_rows, n_cols = arr.shape
    count = 0
    for x in [-1, 0, 1]:
        for y in [-1, 0, 1]:
            if x == 0 and y == 0:
                continue
            if arr[(i + x) % n_rows][(j + y) % n_cols] != arr[i][j]:
                count += 1
    return count

#晶粒尺寸分布统计
# 定义函数确定晶粒边界
def find_grain_boundary(arr):
    rows, cols = arr.shape
    grain_boundary = np.zeros_like(arr)
    for i in range(rows):
        for j in range(cols):
            if (arr[i,j] != arr[(i-1)%rows,(j-1)%cols] or
                arr[i,j] != arr[(i-1)%rows,j] or
                arr[i,j] != arr[(i-1)%rows,(j+1)%cols] or
                arr[i,j] != arr[i,(j-1)%cols] or
                arr[i,j] != arr[i,(j+1)%cols] or
                arr[i,j] != arr[(i+1)%rows,(j-1)%cols] or
                arr[i,j] != arr[(i+1)%rows,j] or
                arr[i,j] != arr[(i+1)%rows,(j+1)%cols]):
                grain_boundary[i,j] = 1
    return grain_boundary

# 定义函数标记晶粒
def label_grain(arr, grain_label, x, y, current_label):
    grain_label[x, y] = current_label
    rows, cols = arr.shape
    neighbors = [(x-1,y-1), (x-1,y), (x-1,y+1),
                 (x,y-1), (x,y+1),
                 (x+1,y-1), (x+1,y), (x+1,y+1)]
    for neighbor in neighbors:
        i, j = neighbor
        if grain_label[i%rows, j%cols] == 0 and arr[i%rows, j%cols] == arr[x, y]:
            label_grain(arr, grain_label, i%rows, j%cols, current_label)

# 定义函数统计晶粒尺寸分布
def count_grain_size(arr):
    rows, cols = arr.shape
    grain_label = np.zeros_like(arr)
    grain_size = {}
    current_label = 0
    for i in range(rows):
        for j in range(cols):
            if grain_label[i,j] == 0:
                current_label += 1
                label_grain(arr, grain_label, i, j, current_label)
                size = np.sum(grain_label == current_label)
                if size in grain_size:
                    grain_size[size] += 1
                else:
                    grain_size[size] = 1
    return grain_size

# 可视化晶粒尺寸分布
def plot_grain_size_distribution(grain_size):
    plt.bar(grain_size.keys(), grain_size.values())
    plt.xlabel('Grain size')
    plt.ylabel('Count')
    plt.title('Grain size distribution')
    plt.show()

# 运行模拟
CA_model(N, n_iter)

